# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/6 11:09 
@Author  : ZhangShenao 
@File    : 5.混合检索.py 
@Desc    : 使用EnsembleRetriever,集成多种检索器,实现混合检索

LangChain内置了EnsembleRetriever集成检索器
它可以接收一个Retriever列表和权重列表作为参数
并根据RRF算法对每个检索器的get_relevant_documents()方法产生的文档列表,按照指定的权重进行排序和融合

集成检索器可以利用不同算法的优势,从而获得比任何单一算法更好的性能
集成检索器一个常见的案例是:将稀疏检索器(如BM25)和密集检索器(如嵌入相似度)结合起来,因为它们的优势是互补的
这种检索方式也被称为混合检索(关键词检索+语义相似度检索)
"""
import dotenv
from langchain.retrievers import EnsembleRetriever
from langchain_community.document_loaders import UnstructuredMarkdownLoader
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.retrievers import BM25Retriever
from langchain_pinecone import PineconeVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter

dotenv.load_dotenv()

# 加载文档列表
loader = UnstructuredMarkdownLoader(file_path="./docs/项目API资料.md", mode="single")
docs = loader.load()
# 使用RecursiveCharacterTextSplitter,对文档进行分割
separators = [
    '\n\n',
    '\n',
    r'。|！|？',
    r'\.\s|\!\s|\?\s',  # 英文标点符号后面通常需要加空格
    r'；|;\s',
    r'，|,\s',
    ' ',
    ''
]
splitter = RecursiveCharacterTextSplitter(
    separators=separators,  # 指定自定义分隔符列表
    chunk_size=100,  # 指定文档分块大小
    chunk_overlap=20,  # 指定文档块之间的重叠部分大小
    add_start_index=True,  # 添加文档块的起始索引
    is_separator_regex=True,  # 指定分隔符为正则表达式形式
)
chunks = splitter.split_documents(docs)

# 基于切割后的文档列表,创建BM25关键词检索器
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 3

# 创建Pinecone向量数据库检索器
embeddings = ZhipuAIEmbeddings(model="embedding-3")
vector_store = PineconeVectorStore(
    index_name="rag",  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace="llm-ops",  # 指定Namespace
    text_key="text",  # 指定文本属性key
)
# 采用相似度检索,并指定相似度得分阈值
vector_store_retriever = vector_store.as_retriever(
    search_type="similarity_score_threshold",
    search_kwargs={
        "k": 3,
        "score_threshold": 0.5,
    }
)

# 创建EnsembleRetriever集成检索器,整合BM25关键词检索和Pinecone语义相似度检索
# 并指定权重列表
retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_store_retriever],
    weights=[0.5, 0.5],
)

# 执行混合检索
results = retriever.invoke("获取应用信息的接口有哪些？")
print(len(results))
print(results[0])
